package main

import (
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"net"
	"net/http"
	"sync"
	"time"
)

// Response 定义了客户端响应的结构
type Response struct {
	Code  int    `json:"code"`  // 状态码
	Error string `json:"error"` // 错误信息（如果有）
	Data  []byte `json:"data"`  // 原始图像数据
}

var (
	clients           = make(map[net.Conn]bool)
	clientsMutex      sync.Mutex
	responseCollector = make(chan string, 1000) // 增加缓冲区大小
)

// sendWithLengthPrefix 使用长度前缀协议发送数据
func sendWithLengthPrefix(conn net.Conn, data []byte) error {
	var lengthBuf [4]byte
	binary.BigEndian.PutUint32(lengthBuf[:], uint32(len(data)))

	// 发送长度前缀
	_, err := conn.Write(lengthBuf[:])
	if err != nil {
		return err
	}

	// 发送实际数据
	_, err = conn.Write(data)
	return err
}

// readWithLengthPrefix 读取长度前缀协议的数据
func readWithLengthPrefix(conn net.Conn) ([]byte, error) {
	var lengthBuf [4]byte
	_, err := io.ReadFull(conn, lengthBuf[:])
	if err != nil {
		return nil, err
	}
	length := binary.BigEndian.Uint32(lengthBuf[:])

	data := make([]byte, length)
	_, err = io.ReadFull(conn, data)
	if err != nil {
		return nil, err
	}

	return data, nil
}

func main() {
	go startTCPServer()
	startHTTPServer()
}

func startTCPServer() {
	listener, err := net.Listen("tcp", ":12345")
	if err != nil {
		fmt.Println("Error listening:", err.Error())
		return
	}
	defer listener.Close()
	fmt.Println("TCP Server listening on :12345")

	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println("Error accepting:", err.Error())
			continue
		}
		fmt.Println("TCP client connected:", conn.RemoteAddr().String())

		clientsMutex.Lock()
		clients[conn] = true
		clientsMutex.Unlock()

		go handleTCPClient(conn)
	}
}

func handleTCPClient(conn net.Conn) {
	defer func() {
		conn.Close()
		clientsMutex.Lock()
		delete(clients, conn)
		clientsMutex.Unlock()
	}()

	for {
		// 读取长度前缀（4字节，大端）
		dataBytes, err := readWithLengthPrefix(conn)
		if err != nil {
			if err == io.EOF {
				fmt.Println("TCP client disconnected:", conn.RemoteAddr().String())
			} else {
				fmt.Println("Error reading data from client:", err)
			}
			return
		}

		// 解析 JSON 数据
		var responseObj Response
		err = json.Unmarshal(dataBytes, &responseObj)
		if err != nil {
			fmt.Println("Error unmarshalling JSON from client:", err)
			continue
		}

		// 将图像数据进行 Base64 编码
		base64Str := base64.StdEncoding.EncodeToString(responseObj.Data)
		fmt.Printf("Received image from %s, Base64 size: %d\n", conn.RemoteAddr().String(), len(base64Str))

		// 这里您可以根据需求进一步处理 base64Str，例如保存到文件或存储到数据库
		// 下面仅仅是打印前100个字符作为示例
		if len(base64Str) > 100 {
			fmt.Println("Base64 Image Data (truncated):", base64Str[:100]+"...")
		} else {
			fmt.Println("Base64 Image Data:", base64Str)
		}

		// 将完整的 Base64 字符串发送到 responseCollector
		responseCollector <- base64Str
	}
}

func sendCaptureCommandToClients() {
	clientsMutex.Lock()
	defer clientsMutex.Unlock()

	messageType := "1" // 消息类型 1 代表截图
	for conn := range clients {
		go func(c net.Conn) {
			err := sendWithLengthPrefix(c, []byte(messageType))
			if err != nil {
				fmt.Printf("Failed to send command to client %s: %v\n", c.RemoteAddr().String(), err)
			} else {
				fmt.Printf("Sent command to client %s\n", c.RemoteAddr().String())
			}
		}(conn)
	}
}

func startHTTPServer() {
	http.HandleFunc("/one", func(w http.ResponseWriter, r *http.Request) {
		clientsMutex.Lock()
		numClients := len(clients)
		clientsMutex.Unlock()

		if numClients == 0 {
			http.Error(w, "No connected clients", http.StatusBadRequest)
			return
		}

		sendCaptureCommandToClients()

		// 等待所有客户端的响应
		var allResponses []string
		var mu sync.Mutex
		var wg sync.WaitGroup

		wg.Add(numClients)
		for i := 0; i < numClients; i++ {
			go func() {
				defer wg.Done()
				select {
				case responseStr := <-responseCollector:
					mu.Lock()
					allResponses = append(allResponses, responseStr)
					mu.Unlock()
				case <-time.After(10 * time.Second):
					fmt.Println("Timeout waiting for client response")
				}
			}()
		}

		wg.Wait()

		// 构建 HTML 模板
		const htmlTemplateStr = `
		<!DOCTYPE html>
		<html lang="zh-CN">
		<head>
			<meta charset="UTF-8">
			<title>截图结果</title>
		</head>
		<body>
			<h1>屏幕截图</h1>
			{{range .Images}}
				<div>
					<img src="data:image/png;base64,{{.}}" alt="Screenshot" style="max-width: 800px; height: auto;" />
				</div>
			{{end}}
		</body>
		</html>
		`

		type PageData struct {
			Images []string
		}

		data := PageData{
			Images: allResponses,
		}

		// 解析并执行模板
		tmpl, err := template.New("result").Parse(htmlTemplateStr)
		if err != nil {
			http.Error(w, "Internal Server Error: unable to parse template", http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		err = tmpl.Execute(w, data)
		if err != nil {
			http.Error(w, "Internal Server Error: unable to execute template", http.StatusInternalServerError)
			return
		}
	})

	fmt.Println("HTTP Server listening on :8848")
	err := http.ListenAndServe(":8848", nil)
	if err != nil {
		fmt.Printf("Failed to start server: %v\n", err)
	}
}
